from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import argparse
import torch

# 使用说明：
# 1. 支持通过命令行参数指定模型路径
# 2. 支持流式输出生成内容
# 示例命令：
# python demo4transformers_stream.py --model_path /path/to/model
# 默认模型路径：/data/models/llm/models/QwQ-32B

# 初始化参数解析器
parser = argparse.ArgumentParser(description='Qwen语言模型推理脚本')
parser.add_argument('--model_path', type=str, 
                    default='/data/models/llm/models/QwQ-32B',
                    help='模型路径（默认为官方32B模型路径）')
args = parser.parse_args()

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    args.model_path,
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(args.model_path)

# 初始化流式输出器
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# 构建对话提示
prompt = "How many r's are in the word \"strawberry\""
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

# 编码输入并生成
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 使用流式输出生成结果
print("[生成开始] 实时输出：")
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32768,
    streamer=streamer,
)

# 最终完整结果解码（可选）
print("\n[生成完成] 最终结果：")
response = tokenizer.decode(generated_ids[0][len(model_inputs.input_ids[0]):], 
                           skip_special_tokens=True)
print(response)